// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/functions/prestosql/aggregates/max_size_for_stats_aggregate.h>
#include <pollux/exec/aggregate.h>
#include <pollux/expression/function_signature.h>
#include <pollux/functions/lib/aggregates/simple_numeric_aggregate.h>
#include <pollux/functions/prestosql/aggregates/aggregate_names.h>
#include <pollux/serializers/presto_serializer.h>

using namespace kumo::pollux::functions::aggregate;

namespace kumo::pollux::aggregate::prestosql {
    namespace {
        std::unique_ptr<VectorSerde> &getVectorSerde() {
            static std::unique_ptr<VectorSerde> serde =
                    std::make_unique<serializer::presto::PrestoVectorSerde>();
            return serde;
        }

        class MaxSizeForStatsAggregate
                : public SimpleNumericAggregate<int64_t, int64_t, int64_t> {
            using BaseAggregate = SimpleNumericAggregate<int64_t, int64_t, int64_t>;

        private:
            std::vector<vector_size_t> elementSizes_;
            std::vector<vector_size_t *> elementSizePtrs_;
            std::vector<IndexRange> elementIndices_;
            DecodedVector decoded_;

        public:
            explicit MaxSizeForStatsAggregate(TypePtr resultType)
                : BaseAggregate(resultType) {
            }

            int32_t accumulatorFixedWidthSize() const override {
                return sizeof(int64_t);
            }

            void extractValues(char **groups, int32_t numGroups, VectorPtr *result)
            override {
                BaseAggregate::doExtractValues(groups, numGroups, result, [&](char *group) {
                    return *BaseAggregate::Aggregate::template value<int64_t>(group);
                });
            }

            void extractAccumulators(char **groups, int32_t numGroups, VectorPtr *result)
            override {
                // Partial and final aggregations are the same.
                extractValues(groups, numGroups, result);
            }

            void addIntermediateResults(
                char **groups,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool mayPushdown) override {
                BaseAggregate::template updateGroups<true, int64_t>(
                    groups,
                    rows,
                    args[0],
                    [](int64_t &result, int64_t value) {
                        if (result < value) {
                            result = value;
                        }
                    },
                    mayPushdown);
            }

            void addSingleGroupIntermediateResults(
                char *group,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool mayPushdown) override {
                BaseAggregate::updateOneGroup(
                    group,
                    rows,
                    args[0],
                    [](int64_t &result, int64_t value) {
                        result = std::max(result, value);
                    },
                    [](int64_t &result, int64_t value, int /* unused */) {
                        result = value;
                    },
                    mayPushdown,
                    static_cast<int64_t>(0));
            }

            void addSingleGroupRawInput(
                char *group,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool /*mayPushdown*/) override {
                doUpdateSingleGroup(group, rows, args[0]);
            }

            void addRawInput(
                char **groups,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool /*mayPushdown*/) override {
                doUpdate(groups, rows, args[0]);
            }

        protected:
            void
            updateOneAccumulator(char *const group, vector_size_t row, int64_t rowSize) {
                if (decoded_.is_null_at(row)) {
                    return;
                }

                // Clear null.
                clearNull(group);
                // Set max(current, this).
                int64_t &current = *value<int64_t>(group);
                current = std::max(current, rowSize);
            }

            void
            doUpdate(char **groups, const SelectivityVector &rows, const VectorPtr &arg) {
                decoded_.decode(*arg, rows, true);

                if (decoded_.isConstantMapping()) {
                    if (!decoded_.is_null_at(0)) {
                        estimateSerializedSizes(arg, rows, 1);
                        rows.applyToSelected([&](auto row) {
                            updateOneAccumulator(groups[row], row, elementSizes_[0]);
                        });
                    }
                } else {
                    estimateSerializedSizes(arg, rows, rows.countSelected());
                    vector_size_t sizeIndex = 0;
                    rows.applyToSelected([&](auto row) {
                        updateOneAccumulator(groups[row], row, elementSizes_[sizeIndex++]);
                    });
                }
            }

            // Estimate the sizes of first numToProcess selected elements in vector.
            void estimateSerializedSizes(
                VectorPtr vector,
                const SelectivityVector &rows,
                vector_size_t numToProcess) {
                elementSizes_.resize(numToProcess);
                std::fill(elementSizes_.begin(), elementSizes_.end(), 0);
                elementIndices_.resize(numToProcess);
                elementSizePtrs_.resize(numToProcess);

                vector_size_t i = 0;
                rows.testSelected([&](auto row) {
                    elementIndices_[i] = IndexRange{row, 1};
                    elementSizePtrs_[i] = &elementSizes_[i];
                    return ++i < numToProcess;
                });

                getVectorSerde()->estimateSerializedSize(
                    vector.get(),
                    melon::Range(elementIndices_.data(), elementIndices_.size()),
                    elementSizePtrs_.data());
            }

            void doUpdateSingleGroup(
                char *group,
                const SelectivityVector &rows,
                const VectorPtr &arg) {
                decoded_.decode(*arg, rows, true);

                if (decoded_.isConstantMapping()) {
                    if (!decoded_.is_null_at(0)) {
                        // Estimate first element because it is constant mapping.
                        estimateSerializedSizes(arg, rows, 1);
                        updateOneAccumulator(group, 0, elementSizes_[0]);
                    }
                } else {
                    estimateSerializedSizes(arg, rows, rows.countSelected());
                    vector_size_t sizeIndex = 0;
                    rows.applyToSelected([&](auto row) {
                        updateOneAccumulator(group, row, elementSizes_[sizeIndex++]);
                    });
                }
            }

            void initializeNewGroupsInternal(
                char **groups,
                melon::Range<const vector_size_t *> indices) override {
                exec::Aggregate::setAllNulls(groups, indices);
                for (auto i: indices) {
                    *BaseAggregate::value<int64_t>(groups[i]) = 0;
                }
            }
        };
    } // namespace

    void registerMaxDataSizeForStatsAggregate(
        const std::string &prefix,
        bool withCompanionFunctions,
        bool overwrite) {
        std::vector<std::shared_ptr<exec::AggregateFunctionSignature> > signatures;

        signatures.push_back(exec::AggregateFunctionSignatureBuilder()
            .typeVariable("T")
            .returnType("bigint")
            .intermediateType("bigint")
            .argumentType("T")
            .build());

        auto name = prefix + kMaxSizeForStats;
        exec::registerAggregateFunction(
            name,
            std::move(signatures),
            [name](
        core::AggregationNode::Step step,
        const std::vector<TypePtr> &argTypes,
        const TypePtr &resultType,
        const core::QueryConfig & /*config*/)
        -> std::unique_ptr<exec::Aggregate> {
                POLLUX_CHECK_EQ(argTypes.size(), 1, "{} takes only one argument", name);
                auto inputType = argTypes[0];

                return std::make_unique<MaxSizeForStatsAggregate>(resultType);
            },
            withCompanionFunctions,
            overwrite);
    }
} // namespace kumo::pollux::aggregate::prestosql
